import numpy as np
import cv2
import os
import sys
import logging
sys.path.append(os.getcwd())
from sku.api import Detector
import sku.api as cpr
from PIL import Image,ImageDraw,ImageFont

labels = ['sku']
colors = [(0,255,0)]

def draw_dets(img, dets,recog):
    font = ImageFont.truetype('sku/data/kaiti.ttf',20)
    imgs = []
    namesdict = {}
    with open('sku/data/folder2classinfo.txt') as f:
        datas = f.readlines()
        for data in datas:
            data = data.strip().split(':')
            namesdict[data[0]] = data[1]
    tl = round(0.002 * (img.shape[0] + img.shape[1]) / 2) + 1
    tf = max(tl - 1, 1)  # font thickness
    if dets[0] is not None:
        for index,det in enumerate(dets):
            try:
                cls, score, bbox = int(det[5]), det[4], det[:4]
            except:
                print()
            color = colors[cls]
            c1, c2 = (int(bbox[0]), int(bbox[1])), (int(bbox[2]), int(bbox[3]))
            imgs.append(img[c1[1]:c2[1],c1[0]:c2[0]].copy())
            cv2.rectangle(img, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA)
            cv2.putText(img, str(index), ((c1[0]+c2[0])//2, (c1[1]+c2[1])//2), 0, tl / 4, [0,0, 255], thickness=tf, lineType=cv2.LINE_AA)

    img_black = np.ones_like(img) * 255
    img_black = np.rot90(img_black)
    img_black = Image.fromarray(img_black)
    draw = ImageDraw.Draw(img_black)
    imgs = recog.preprocess(imgs)
    feat_vec = recog.inference(imgs)
    ids, scores = recog.classify(feat_vec)
    unreliable_Ids = np.where(scores < recog.thr)[0]
    ids[unreliable_Ids] = 0
    x = 0
    y = 0
    for index,id in enumerate(ids):
        draw.text((x, y),str(index) + ':' + namesdict[str(id)],(0,0,0),font=font)
        y += 40
        if y >= img.shape[1]:
            y = 0
            x = img.shape[0] // 2
    img_black = np.array(img_black)
    return img,img_black


def find_files(folder,ext):
    g = os.walk(folder)
    L=[]
    for path,dir_list,file_list in g:  
        for file_name in file_list:  
            if os.path.splitext(file_name)[1] == ext:
                L.append(file_name)
    return L

def predict():
    logging.basicConfig(level=logging.DEBUG)
    model_name = 'sku_yolov5_0.0.1.onnx'
    det = Detector(model_path=os.path.join('sku/data', model_name), batch_size=1)
    recog = cpr.SkuRecognizer('sku/data/sku_recog_mobilenetV2_0.0.2.onnx', mode=cpr.ie.MODE_ORT, allow_growth=True,
                              per_process_gpu_memory_fraction=None, batch_size=16)

    recog.load_templates('sku/data/sku_feats_0.0.2.npy', 'sku/data/sku_ids_0.0.2.npy')
    addr = '/home/lixuan/workspace/project/CatVsDog/media/img/'
    img_names = find_files(addr,'.jpg')
    imgs = [cv2.imread(addr+img_name) for img_name in img_names]
    outputs = det.detect(imgs)
    for img,output,img_name in zip(imgs,outputs,img_names):
        img,img_black = draw_dets(img, output,recog)
        img = cv2.resize(img,(900,600))
        img_black = cv2.resize(img_black, (600,900))
        cv2.imwrite('/home/lixuan/workspace/project/CatVsDog/media/img/test1.jpg',img)
        cv2.imwrite('/home/lixuan/workspace/project/CatVsDog/media/img/test2.jpg',img_black)